from pyalgotrade.barfeed.csvfeed import GenericBarFeed
from pyalgotrade.bar import Frequency

from pyalgotrade.technical import ma
from pyalgotrade.technical import rsi
from pyalgotrade.technical import cross

from pyalgotrade import strategy
from pyalgotrade.stratanalyzer import returns
from pyalgotrade.stratanalyzer import sharpe
from pyalgotrade.stratanalyzer import drawdown
from pyalgotrade.stratanalyzer import trades
from pyalgotrade import plotter

class MyRSIStrategy(strategy.BacktestingStrategy):
    def __init__(self, feed, instrument, entrySMA, exitSMA, rsiPeriod, overBoughtThreshold, overSoldThreshold):
        super(MyRSIStrategy, self).__init__(feed)
        self.__instrument = instrument
        # We'll use adjusted close values, if available, instead of regular close values.
        if feed.barsHaveAdjClose():
            self.setUseAdjustedValues(True)
        self.__priceDS = feed[instrument].getPriceDataSeries()
        self.__entrySMA = ma.SMA(self.__priceDS, entrySMA)
        self.__exitSMA = ma.SMA(self.__priceDS, exitSMA)
        self.__rsi = rsi.RSI(self.__priceDS, rsiPeriod)
        self.__overBoughtThreshold = overBoughtThreshold
        self.__overSoldThreshold = overSoldThreshold
        self.__longPos = None

    def getEntrySMA(self):
        return self.__entrySMA

    def getExitSMA(self):
        return self.__exitSMA

    def getRSI(self):
        return self.__rsi

    def onEnterCanceled(self, position):
        if self.__longPos == position:
            self.__longPos = None

    def onExitOk(self, position):
        if self.__longPos == position:
            self.__longPos = None

    def onExitCanceled(self, position):
        # If the exit was canceled, re-submit it.
        position.exitMarket()

    def onBars(self, bars):
        # Wait for enough bars to be available to calculate SMA and RSI.
        if self.__exitSMA[-1] is None or self.__entrySMA[-1] is None or self.__rsi[-1] is None:
            return

        bar = bars[self.__instrument]
        if self.__longPos is not None:
            if self.exitLongSignal():
                self.__longPos.exitMarket()
        else:
            if self.enterLongSignal(bar):
                shares = int(self.getBroker().getCash() * 0.9 / bars[self.__instrument].getPrice())
                self.__longPos = self.enterLong(self.__instrument, shares, True)

    def enterLongSignal(self, bar):
        return bar.getPrice() > self.__entrySMA[-1] and self.__rsi[-1] <= self.__overSoldThreshold

    def exitLongSignal(self):
        return cross.cross_above(self.__priceDS, self.__exitSMA) and not self.__longPos.exitActive()

'''
instrument="Sinopec"
entrySMA=163
exitSMA=7
rsiPeriod=2
overBoughtThreshold=90
overSoldThreshold=10

feed=GenericBarFeed(Frequency.DAY,None,None)
feed.addBarsFromCSV("Sinopec",r"c:\Users\26356\Desktop\python\BUY\600028SH.new.csv")

myRSIStrategy=MyRSIStrategy(feed,instrument,entrySMA,exitSMA,rsiPeriod,overBoughtThreshold,overSoldThreshold)

''''''
returnsAnalyzer=returns.Returns()
myRSIStrategy.attachAnalyzer(returnsAnalyzer)
sharpeRatioAnalyzer=sharpe.SharpeRatio()
myRSIStrategy.attachAnalyzer(sharpeRatioAnalyzer)
drawdownAnalyzer=drawdown.DrawDown()
myRSIStrategy.attachAnalyzer(drawdownAnalyzer)
tradesAnalyzer=trades.Trades()
myRSIStrategy.attachAnalyzer(tradesAnalyzer)

''''''
plt=plotter.StrategyPlotter(myRSIStrategy)
plt.getInstrumentSubplot(instrument).addDataSeries("ENtry SMA",myRSIStrategy.getEntrySMA())
plt.getInstrumentSubplot(instrument).addDataSeries("Exit SMA",myRSIStrategy.getExitSMA())
plt.getOrCreateSubplot("rsi").addDataSeries("RSI",myRSIStrategy.getRSI())
plt.getOrCreateSubplot("rsi").addLine("Overbought",overBoughtThreshold)
plt.getOrCreateSubplot("rsi").addLine("Oversold",overSoldThreshold)
plt.getOrCreateSubplot("returns").addDataSeries("Simple returns",returnsAnalyzer.getReturns())

''''''
myRSIStrategy.run()
print("Final portfolio value1:$%.2f"%(myRSIStrategy.getBroker().getEquity()))
print("Final portfolio value2:$%.2f"%(myRSIStrategy.getResult()))
print("Cumulative returns:%.2f%%"%(returnsAnalyzer.getCumulativeReturns()[-1]*100))
print("Sharpe ratio:%.2f"%(sharpeRatioAnalyzer.getSharpeRatio(0.03)))
print("Max.drawdown:%.2f%%"%(drawdownAnalyzer.getMaxDrawDown()*100))
print("Longest drawdown duration:%s"%(drawdownAnalyzer.getLongestDrawDownDuration()))
plt.plot()
'''